/**
 * ws client reconnect example
 * Reconect handle:
 * - WEBSOCKET_EVENT_FINISH
 * - WEBSOCKET_EVENT_DISCONNECTED
 * - WEBSOCKET_EVENT_CLOSED
 * - WEBSOCKET_EVENT_ERROR
 */

#include <stdio.h>
#include "esp_wifi.h"
#include "esp_system.h"
#include "nvs_flash.h"
#include "esp_event.h"
#include "protocol_examples_common.h"
#include "esp_crt_bundle.h"

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"
#include "freertos/event_groups.h"

#include "esp_log.h"
#include "esp_websocket_client.h"
#include "esp_event.h"
#include <cJSON.h>

#include "ws_client.h"

static const char *TAG = "ws";

#define RECONNECT_DELAY_MS 5000
#define RECONNECT_BIT (1 << 0)

typedef struct
{
    uint16_t len;
    uint8_t data[WS_DATA_BUFFER_SIZE];
} ws_client_data_t;

static esp_websocket_client_handle_t ws_client = NULL;
static QueueHandle_t ws_client_tx_queue = NULL;
static EventGroupHandle_t ws_event_group;
static ws_data_callback_t ws_data_callback = NULL;
static bool is_destroy = false;

static void log_error_if_nonzero(const char *message, int error_code)
{
    if (error_code != 0)
    {
        ESP_LOGE(TAG, "Last error %s: 0x%x", message, error_code);
    }
}

static void websocket_event_handler(void *handler_args, esp_event_base_t base, int32_t event_id, void *event_data)
{
    esp_websocket_client_handle_t ws_client = (esp_websocket_client_handle_t)handler_args;
    esp_websocket_event_data_t *data = (esp_websocket_event_data_t *)event_data;
    switch (event_id)
    {
        case WEBSOCKET_EVENT_BEGIN:
            ESP_LOGI(TAG, "WEBSOCKET_EVENT_BEGIN");
            break;

        case WEBSOCKET_EVENT_CONNECTED:
            ESP_LOGI(TAG, "WEBSOCKET_EVENT_CONNECTED");
            is_destroy = false;
            break;

        case WEBSOCKET_EVENT_DISCONNECTED:
            ESP_LOGI(TAG, "WEBSOCKET_EVENT_DISCONNECTED");
            log_error_if_nonzero("HTTP status code", data->error_handle.esp_ws_handshake_status_code);
            if (data->error_handle.error_type == WEBSOCKET_ERROR_TYPE_TCP_TRANSPORT)
            {
                log_error_if_nonzero("reported from esp-tls", data->error_handle.esp_tls_last_esp_err);
                log_error_if_nonzero("reported from tls stack", data->error_handle.esp_tls_stack_err);
                log_error_if_nonzero("captured as transport's socket errno", data->error_handle.esp_transport_sock_errno);
            }
            esp_websocket_client_set_reconnect_timeout(ws_client, 5000);
            break;

        case WEBSOCKET_EVENT_DATA:
            ESP_LOGI(TAG, "WEBSOCKET_EVENT_DATA");
            ESP_LOGI(TAG, "Received opcode=%d", data->op_code);
            if (data->op_code == 0x1 || data->op_code == 0x2)
            {
                if (ws_data_callback != NULL)
                {
                    ws_data_callback(data->op_code, (const uint8_t *)data->data_ptr, data->data_len);
                }
            }
            else if (data->op_code == 0x08 && data->data_len == 2)
            {
                ESP_LOGW(TAG, "Received closed message with code=%d", 256 * data->data_ptr[0] + data->data_ptr[1]);
            }
            else if (data->op_code == 0x9)
            {
                ESP_LOGD(TAG, "WS PING.");
            }
            else if (data->op_code == 0x0A)
            {
                ESP_LOGD(TAG, "WS PONG.");
            }
            else
            {
                ESP_LOGW(TAG, "Received=%.*s\n\n", data->data_len, (char *)data->data_ptr);
            }
            break;

        case WEBSOCKET_EVENT_ERROR:
            ESP_LOGI(TAG, "WEBSOCKET_EVENT_ERROR");
            log_error_if_nonzero("HTTP status code", data->error_handle.esp_ws_handshake_status_code);
            if (data->error_handle.error_type == WEBSOCKET_ERROR_TYPE_TCP_TRANSPORT)
            {
                log_error_if_nonzero("reported from esp-tls", data->error_handle.esp_tls_last_esp_err);
                log_error_if_nonzero("reported from tls stack", data->error_handle.esp_tls_stack_err);
                log_error_if_nonzero("captured as transport's socket errno", data->error_handle.esp_transport_sock_errno);
            }
            break;

        case WEBSOCKET_EVENT_FINISH:
            ESP_LOGI(TAG, "WEBSOCKET_EVENT_FINISH");
            is_destroy = true;
            xEventGroupSetBits(ws_event_group, RECONNECT_BIT);
            break;

        case WEBSOCKET_EVENT_CLOSED:
            ESP_LOGI(TAG, "WEBSOCKET_EVENT_CLOSED");
            break;
    }
}

static void ws_app_start(void)
{
    esp_websocket_client_config_t websocket_cfg = {};
    websocket_cfg.uri = CONFIG_WEBSOCKET_URI;
    websocket_cfg.reconnect_timeout_ms = 2000;

    ESP_LOGI(TAG, "Connecting to %s...", websocket_cfg.uri);
    ws_client = esp_websocket_client_init(&websocket_cfg);
    esp_websocket_register_events(ws_client, WEBSOCKET_EVENT_ANY, websocket_event_handler, (void *)ws_client);

    esp_websocket_client_start(ws_client);
}

static void ws_reconnect_task(void *arg)
{
    while (1)
    {
        EventBits_t bits = xEventGroupWaitBits(ws_event_group, RECONNECT_BIT, pdTRUE, pdFALSE, portMAX_DELAY);
        if (bits & RECONNECT_BIT)
        {
            if (ws_client != NULL)
            {
                esp_websocket_client_close(ws_client, portMAX_DELAY);
                ESP_LOGI(TAG, "Websocket Stopped");
                // make sure to add delay here
                vTaskDelay(pdMS_TO_TICKS(1000));
                esp_websocket_client_destroy(ws_client);
                ws_client = NULL;

                printf("start to reconnect...\n");
                vTaskDelay(pdMS_TO_TICKS(RECONNECT_DELAY_MS));
                ws_app_start();
            }
        }
    }
    vTaskDelete(NULL);
}

static void ws_client_tx_task(void *arg)
{
    ws_client_data_t tx;
    for (;;)
    {
        if (xQueueReceive(ws_client_tx_queue, &tx, portMAX_DELAY) == pdTRUE)
        {
            if (esp_websocket_client_is_connected(ws_client) && !is_destroy)
            {
                esp_websocket_client_send_text(ws_client, (char *)tx.data, tx.len, portMAX_DELAY);
            }
            else
            {
                ESP_LOGW(TAG, "ws client is disconnected.");
            }
        }
    }
    vTaskDelete(NULL);
}

void ws_client_setup(void)
{
    ws_event_group = xEventGroupCreate();
    ws_client_tx_queue = xQueueCreate(10, sizeof(ws_client_data_t));
}

int ws_client_start(void)
{
    ws_app_start();

    xTaskCreate(ws_client_tx_task, "ws_client_tx_task", 8192, NULL, 8, NULL);
    xTaskCreate(ws_reconnect_task, "ws_reconnect_task", 4096, NULL, 6, NULL);
    return 0;
}

int ws_client_send(uint8_t *data, uint16_t len)
{
    ws_client_data_t tx = {0};
    memcpy(tx.data, data, len);
    tx.len = len;
    if (xQueueSend(ws_client_tx_queue, &tx, pdMS_TO_TICKS(50)) != pdTRUE)
    {
        ESP_LOGW(TAG, "ws_client_tx_queue is full");
    }

    return 0;
}

void ws_set_data_callback(ws_data_callback_t callback)
{
    ws_data_callback = callback;
}
